from PyQt5 import QtWidgets, QtCore

from demo.untitled import Ui_Form
import sys
import cv2
import numpy as np
from PyQt5.QtCore import QThread, pyqtSignal, QTimer
from PyQt5.QtGui import QImage, QPixmap
from PyQt5.QtWidgets import QApplication, QWidget, QMessageBox, QFileDialog
from utils.config import Config
from utils.Camera import CameraClass
import time


class CameraCalibrationThread(QThread):
    finish_signal = pyqtSignal(bool, str)

    def __init__(self):
        super(CameraCalibrationThread, self).__init__()
        Config.camera_form = Config.camera_form

    def run(self):
        if Config.camera_form == '单目':
            _, cameraMatrix, distCoeffs, rvecs, tvecs = cv2.calibrateCamera(Config.objpoints, Config.imgpoints_l,
                                                                            (
                                                                                int(Config.camera_horizontal_resolution / 2),
                                                                                Config.camera_vertical_resolution),
                                                                            None,
                                                                            None)
            with open(Config.save_path, 'w') as f:
                f.write('相机矩阵：\n')
                f.write(np.array2string(cameraMatrix, separator=','))
                f.write('\n畸变系数:\n')
                f.write(np.array2string(distCoeffs, separator=','))
        else:
            _, cameraMatrix_l, distCoeffs_l, rvecs_l, tvecs_L = cv2.calibrateCamera(Config.objpoints,
                                                                                    Config.imgpoints_l,
                                                                                    (
                                                                                        int(Config.camera_horizontal_resolution / 2),
                                                                                        Config.camera_vertical_resolution),
                                                                                    None,
                                                                                    None)
            _, cameraMatrix_r, distCoeffs_r, rvecs_r, tvecs_r = cv2.calibrateCamera(Config.objpoints,
                                                                                    Config.imgpoints_r,
                                                                                    (
                                                                                        int(Config.camera_horizontal_resolution / 2),
                                                                                        Config.camera_vertical_resolution),
                                                                                    None,
                                                                                    None)
            # 进行标定
            _, _, _, _, _, R, T, E, F = cv2.stereoCalibrate(Config.objpoints, Config.imgpoints_l, Config.imgpoints_r,
                                                            cameraMatrix_l, distCoeffs_l, cameraMatrix_r, distCoeffs_r,
                                                            (int(Config.camera_horizontal_resolution / 2),
                                                             Config.camera_vertical_resolution),
                                                            flags=cv2.CALIB_USE_INTRINSIC_GUESS)
            # 计算矫正变换矩阵
            R1, R2, P1, P2, Q, _, _ = cv2.stereoRectify(cameraMatrix_l, distCoeffs_l, cameraMatrix_r, distCoeffs_r,
                                                        (int(Config.camera_horizontal_resolution / 2),
                                                         Config.camera_vertical_resolution), R, T, flags=0)

            # 保存相机参数
            with open(Config.save_path, 'w') as f:
                f.write('左相机矩阵：\n')
                f.write(np.array2string(cameraMatrix_l, separator=','))
                f.write('\n左相机畸变系数:\n')
                f.write(np.array2string(distCoeffs_l, separator=','))
                f.write('\n右相机矩阵：\n')
                f.write(np.array2string(cameraMatrix_r, separator=','))
                f.write('\n右相机畸变系数:\n')
                f.write(np.array2string(distCoeffs_r, separator=','))
                f.write('\n旋转矩阵:\n')
                f.write(np.array2string(R, separator=','))
                f.write('\n平移矩阵:\n')
                f.write(np.array2string(T, separator=','))
                f.write('\n重投影矩阵:\n')
                f.write(np.array2string(Q, separator=','))
            # np.savez(Config.save_path, camera_matrix_l=cameraMatrix_l, dist_coeffs_l=distCoeffs_l,
            #          camera_matrix_r=cameraMatrix_r,
            #          dist_coeffs_r=distCoeffs_r, Q=Q)
        self.finish_signal.emit(True, '标定成功')


class VideoPictureCapture(QThread):
    counter_signal = pyqtSignal()

    def __init__(self):
        super(VideoPictureCapture, self).__init__()

    def run(self):
        # time.sleep(1)
        if Config.camera_form == '单目':
            self.mono_photos_capture()
        else:
            self.stereo_photos_capture()

        # 定义单目图像采集

    def mono_photos_capture(self):
        gray = cv2.cvtColor(Config.frame, cv2.COLOR_BGR2GRAY)
        # 检测棋盘格
        ret, corners = cv2.findChessboardCorners(gray, Config.pattern_size)
        if ret:
            time.sleep(1)
            _, corners = cv2.find4QuadCornerSubpix(gray, corners, (5, 5))
            Config.imgpoints_l.append(corners)
            objp = np.zeros((Config.pattern_size[0] * Config.pattern_size[1], 3), np.float32)
            objp[:, :2] = np.mgrid[0:Config.pattern_size[0], 0:Config.pattern_size[1]].T.reshape(-1,
                                                                                                 2) * Config.square_size
            Config.objpoints.append(objp)

            # 更新采集计数器
            Config.capture_counter += 1
            self.counter_signal.emit()

    # 定义双目图像采集
    def stereo_photos_capture(self):

        img_l = Config.frame[:, 0:int(Config.camera_horizontal_resolution / 2), :]
        img_r = Config.frame[:,
                int(Config.camera_horizontal_resolution / 2):Config.camera_horizontal_resolution, :]

        gray_l = cv2.cvtColor(img_l, cv2.COLOR_BGR2GRAY)
        gray_r = cv2.cvtColor(img_r, cv2.COLOR_BGR2GRAY)

        # 检测棋盘格
        ret_l, corners_l = cv2.findChessboardCorners(gray_l, Config.pattern_size)
        ret_r, corners_r = cv2.findChessboardCorners(gray_r, Config.pattern_size)
        if ret_r and ret_l:
            time.sleep(1)
            _, corners_l = cv2.find4QuadCornerSubpix(gray_l, corners_l, (5, 5))
            _, corners_r = cv2.find4QuadCornerSubpix(gray_r, corners_r, (5, 5))

            Config.imgpoints_l.append(corners_l)
            Config.imgpoints_r.append(corners_r)

            objp = np.zeros((Config.pattern_size[0] * Config.pattern_size[1], 3), np.float32)
            objp[:, :2] = np.mgrid[0:Config.pattern_size[0], 0:Config.pattern_size[1]].T.reshape(-1,
                                                                                                 2) * Config.square_size
            Config.objpoints.append(objp)

            # 更新采集计数器
            Config.capture_counter += 1


class CameraCalibrationWindow(QWidget, Ui_Form):

    def __init__(self):
        super(CameraCalibrationWindow, self).__init__()
        # 初始化定时器
        self.qtimer_1 = QTimer()
        # self.qtimer_2 = QTimer()
        self.setupUi(self)
        self.initUi()

    def initUi(self):
        # 初始化各参数
        Config.camera_form = str(self.comboBox_2.currentText())
        Config.camera_horizontal_resolution = int(self.spinBox.value())
        Config.camera_vertical_resolution = int(self.spinBox_2.value())
        Config.capture_total = int(self.spinBox_3.value())
        Config.pattern_size = list(map(int, self.comboBox.currentText().split('x')))
        Config.square_size = int(self.doubleSpinBox.value())
        Config.camera = CameraClass(Config.camera_horizontal_resolution,
                                    Config.camera_vertical_resolution, 0)
        Config.camera.camera_init()

        self.qtimer_1.timeout.connect(self.all_photos_capture)
        # self.qtimer_2.timeout.connect(self.update_counter)

        self.comboBox_2.currentIndexChanged.connect(self.update_form)

        self.pushButton_2.setEnabled(True)
        self.pushButton_3.setEnabled(True)
        self.pushButton_4.setEnabled(False)

    # 定义计数更新槽函数
    def update_counter(self):
        self.label_7.setText(str(Config.capture_counter))

    # 定义相机模式切换槽函数
    def update_form(self):
        Config.camera_form = str(self.comboBox_2.currentText())

    def all_photos_capture(self):
        Config.frame = Config.camera.camera_datas_capture()
        if Config.frame is not None:
            frame = Config.camera.bgr_to_rgb(Config.frame)
            frame = QImage(frame.data, frame.shape[1], frame.shape[0], frame.shape[1] * 3, QImage.Format_RGB888)
            frame = QPixmap.fromImage(frame)
            self.update_parameter()
            # 自适应图片大小
            self.label.setPixmap(frame.scaled(self.label.width(),
                                              self.label.height(),
                                              QtCore.Qt.KeepAspectRatio))
            # print(Config.square_size)
            # print(Config.capture_counter)
            # print(Config.pattern_size)
            # print(Config.capture_total)
            # print(Config.camera_form)
            self.thread_1.start()

    # 定义开始采集槽函数
    def start_capture(self):
        self.qtimer_1.start(35)
        # self.qtimer_2.start(35)

        self.thread_1 = VideoPictureCapture()
        self.pushButton_2.setEnabled(True)
        self.pushButton_3.setEnabled(True)
        self.pushButton_4.setEnabled(True)

    # 定义开始标定槽函数
    def start_calibration(self):
        if Config.capture_counter <= Config.capture_total:
            QMessageBox.warning(self, '错误', '请先采集图像')
            return
        if self.pushButton.text() == '选择路径':
            QMessageBox.warning(self, '错误', '请选择参数保存路径')
            return
        self.qtimer_1.stop()
        self.calibration_thread = CameraCalibrationThread()
        if Config.camera_form == '单目':
            self.calibration_thread.start()
            self.calibration_thread.finish_signal.connect(self.success_calibration)
        else:
            self.calibration_thread.start()
            self.calibration_thread.finish_signal.connect(self.success_calibration)

    # 定义成功标定槽函数
    def success_calibration(self, ret, text):
        if ret:
            self.label_9.setText(text)

    # 定义选择参数保存路径槽函数
    def choice_path(self):
        save_path, _ = QFileDialog.getSaveFileName(self, '保存参数', '', 'Text File (*.txt)')
        if save_path:
            self.pushButton.setText(save_path)
            Config.save_path = save_path

    # 参数初始化刷新
    def refresh(self):
        self.qtimer_1.stop()
        # self.qtimer_2.stop()
        # 初始化各参数
        Config.camera.camera_exit()
        Config.camera_form = None
        Config.objpoints = []
        Config.imgpoints_l = []
        Config.imgpoints_r = []
        Config.capture_counter = 0
        self.thread_1.terminate()
        Config.camera_form = str(self.comboBox_2.currentText())
        Config.camera_horizontal_resolution = int(self.spinBox.value())
        Config.camera_vertical_resolution = int(self.spinBox_2.value())
        Config.capture_total = int(self.spinBox_3.value())
        Config.pattern_size = list(map(int, self.comboBox.currentText().split('x')))
        Config.square_size = int(self.doubleSpinBox.value())

        Config.camera = CameraClass(Config.camera_horizontal_resolution,
                                    Config.camera_vertical_resolution, 0)
        Config.camera.camera_init()
        self.label_9.setText('未标定')
        self.label_7.setText(str(0))
        self.pushButton_2.setEnabled(True)
        self.pushButton_3.setEnabled(True)
        self.pushButton_4.setEnabled(False)

    # 参数实时更新
    def update_parameter(self):
        Config.camera_form = str(self.comboBox_2.currentText())
        Config.camera_horizontal_resolution = int(self.spinBox.value())
        Config.camera_vertical_resolution = int(self.spinBox_2.value())
        Config.capture_total = int(self.spinBox_3.value())
        Config.pattern_size = list(map(int, self.comboBox.currentText().split('x')))
        Config.square_size = int(self.doubleSpinBox.value())
        self.label_7.setText(str(Config.capture_counter))

    def closeEvent(self, event):
        result = QtWidgets.QMessageBox.question(self, "标题", "是否确定关闭！！！",
                                                QtWidgets.QMessageBox.Yes | QtWidgets.QMessageBox.No)
        if result == QtWidgets.QMessageBox.Yes:
            event.accept()
            Config.camera.camera_exit()
        else:
            event.ignore()


if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = CameraCalibrationWindow()
    window.show()
    sys.exit(app.exec_())
